#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@FileName: allJyt.py
@Author: Hugo Wang
@Date: 2025-06-08 11:17
@Project: company_talent_loss
"""
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, precision_score, recall_score
from sklearn.linear_model import LogisticRegression
import pandas as pd
import os
from sklearn.preprocessing import StandardScaler
import joblib

def pre_data(file_path):
    """
    数据预处理
    :return:
    """
    df = pd.read_csv(file_path)
    # 被删除的特征列
    df.drop(
        ['Department', 'EmployeeNumber', 'Gender', 'Over18', 'StandardHours', 'EducationField', 'JobRole'],
        axis='columns', inplace=True)

    df['MaritalStatus'] = df['MaritalStatus'].map({
        'Divorced': 0,
        'Single': 1,
        'Married': 2
    }).fillna(-1).astype(int)

    df['BusinessTravel'] = df['BusinessTravel'].map({
        'Travel_Rarely': 0,
        'Travel_Frequently': 1,
        'Non-Travel': 2
    }).fillna(-1).astype(int)

    # print(df2.info())
    # 去掉object对象产生的含义重复的列
    # df = df2.drop(['Gender_Male', 'OverTime_No'], axis=1)
    df['JobSatisfactionScore'] = (df['JobSatisfaction'] * 0.3 +
                                  df['EnvironmentSatisfaction'] * 0.2 +
                                  df['RelationshipSatisfaction'] * 0.2 +
                                  df['WorkLifeBalance'] * 0.15 +
                                  df['JobInvolvement'] * 0.15)
    # 计算员工稳定性（工作满意度/（呆过公司数量+1））
    df['EmployeeStability'] = (df['JobSatisfaction'] / (df['NumCompaniesWorked'] + 1))
    # 工作压力('WorkPressure') = 是否加班(0verTime)x(5 - 工作生活平衡度WorkLifeBalance)
    df['OverTime_num'] = df['OverTime'].map({'Yes': 1, 'No': 0})

    # 工作压力('WorkPressure') = 是否加班(0verTime)x(5 - 工作生活平衡度WorkLifeBalance)
    df['WorkPressure'] = df['OverTime_num'] * (5 - df['WorkLifeBalance'])
    # 资源满意度(ResourceSatisfaction) = (工作环境满意度 (EnvironmentSatisfaction)+ 人际关系满意度
    df['ResourceSatisfaction'] = (df['EnvironmentSatisfaction'] + df['RelationshipSatisfaction']) / 2

    # 职业停滞(CareerStagnation) = 自上次晋升以来的年限(YearsSinceLastPromotion)/(在公司工作年限
    df['CareerStagnation'] = df['YearsSinceLastPromotion'] / (df['YearsAtCompany'] + 0.001)

    # 分离标签
    x = df.drop('Attrition', axis=1)  # 等价于 df.iloc[:, 1:]，但语义更明确
    y = df['Attrition']

    x = pd.get_dummies(x).drop('OverTime_No', axis=1)

    return x, y


if __name__ == '__main__':
    x_train_org, y_train_org = pre_data(r'../data/train.csv')
    test_x, test_y = pre_data('../data/test2.csv')

    # 训练集
    x_train, x_test, y_train, y_test = train_test_split(x_train_org, y_train_org, test_size=0.2, random_state=35,
                                                        stratify=y_train_org)
    transfer = StandardScaler()
    x_train = transfer.fit_transform(x_train)
    x_test = transfer.transform(x_test)
    test_x = transfer.transform(test_x)

    # 1.逻辑回归
    es = LogisticRegression(fit_intercept=True, random_state=2)
    es.fit(x_train, y_train)
    y_pre = es.predict_proba(x_test)[:, 1]
    print("当前模型 逻辑回归 的训练集roc评分为：{}".format(roc_auc_score(y_test, y_pre)))
    joblib.dump(es, "../model/LOG_20250608.pkl")
    print(f"模型 逻辑回归 保存成功，保存路径{os.path.abspath("../model/LOG_20250608.pkl")}")

    # 测试集测试
    y_pre = es.predict(test_x)
    test_y_pre = es.predict_proba(test_x)[:, 1]
    print("当前模型的测试集roc评分为：{}".format(roc_auc_score(test_y, test_y_pre)))

    #2. 决策树CART
    from sklearn.tree import DecisionTreeClassifier

    es2 = DecisionTreeClassifier(criterion="gini", max_depth=50, random_state=2)
    es2.fit(x_train, y_train)
    y_pre = es2.predict_proba(x_test)[:, 1]
    print("当前模型 决策树CART 的训练集roc评分为：{}".format(roc_auc_score(y_test, y_pre)))
    joblib.dump(es2, "../model/dtcCART_20250608.pkl")
    print(f"决策树CART 模型保存成功，保存路径{os.path.abspath("../model/dtcCART_20250608.pkl")}")
    # 测试集测试
    test_y_pre = es2.predict_proba(test_x)[:, 1]
    print("当前模型  决策树CART 的测试集roc评分为：{}".format(roc_auc_score(test_y, test_y_pre)))

    #3 决策树 ID3.0
    from sklearn.tree import DecisionTreeClassifier

    es3 = DecisionTreeClassifier(criterion="entropy", max_depth=50, random_state=14)
    es3.fit(x_train, y_train)
    y_pre = es3.predict(x_test)
    print("当前模型 决策树 ID3.0的roc评分为：{}".format(roc_auc_score(y_test, y_pre)))

    joblib.dump(es3, "../model/dtcID3_20250608.pkl")
    print(f"模型 决策树 ID3.0 保存成功，保存路径{os.path.abspath("../model/dtcID3_20250608.pkl")}")
    # 测试集测试
    test_y_pre = es3.predict_proba(test_x)[:, 1]
    print("当前模型 决策树 ID3.0 的测试集roc评分为：{}".format(roc_auc_score(test_y, test_y_pre)))

    #4 随机森林ID3.0
    from sklearn.ensemble import RandomForestClassifier

    es4 = RandomForestClassifier(criterion="entropy", max_depth=50, random_state=34)
    es4.fit(x_train, y_train)
    y_pre = es4.predict_proba(x_test)[:, 1]
    print("当前模型 随机森林ID3.0 的训练集roc评分为：{}".format(roc_auc_score(y_test, y_pre)))

    joblib.dump(es4, "../model/rtID3.0_20250608.pkl")
    print(f"模型保存成功，保存路径{os.path.abspath("../model/rtID3.0_20250608.pkl")}")

    # 测试集测试
    test_y_pre = es4.predict_proba(test_x)[:, 1]
    print("随机森林ID3.0 当前模型的测试集roc评分为：{}".format(roc_auc_score(test_y, test_y_pre)))

    #5 随机森林CART
    from sklearn.ensemble import RandomForestClassifier

    es5 = RandomForestClassifier(criterion="gini", max_depth=50, random_state=10)
    es5.fit(x_train, y_train)
    y_pre = es5.predict_proba(x_test)[:, 1]
    print("当前模型 随机森林CART 的训练集roc评分为：{}".format(roc_auc_score(y_test, y_pre)))

    joblib.dump(es5, "../model/rtCART_20250608.pkl")
    print(f" 随机森林CART 模型保存成功，保存路径{os.path.abspath("../model/rtCART_20250608.pkl")}")

    # 测试集测试
    test_y_pre = es.predict_proba(test_x)[:, 1]
    print("当前模型 随机森林CART 的测试集roc评分为：{}".format(roc_auc_score(test_y, test_y_pre)))

    #6 随机森林xgboost
    import xgboost as xgb

    es6 = xgb.XGBClassifier(max_depth=1, gamma=10, n_estimators=201, random_state=50)
    es6.fit(x_train, y_train)
    y_pre = es6.predict_proba(x_test)[:, 1]
    print("当前模型的训练集roc评分为：{}".format(roc_auc_score(y_test, y_pre)))

    joblib.dump(es6, "../model/xgboost_20250608.pkl")
    print(f"模型保存成功，保存路径{os.path.abspath("../model/xgboost_20250608.pkl")}")

    # 测试集测试
    y_pre = es6.predict(test_x)
    test_y_pre = es6.predict_proba(test_x)[:, 1]
    print('随机森林xgboost 精确率为：{}'.format(precision_score(test_y, y_pre)))
    print('随机森林xgboost 召回率为：{}'.format(recall_score(test_y, y_pre)))
    print("当前模型 随机森林xgboost 的测试集roc评分为：{}".format(roc_auc_score(test_y, test_y_pre)))
